import os
import sys
# 获取当前文件夹的绝对路径，然后向上找2层到项目根目录
project_root = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../..')
sys.path.append(project_root)
import torch
from P03_NER.LSTM_CRF.config import Config
from P03_NER.LSTM_CRF.model.BiLSTM import NERLSTM
from P03_NER.LSTM_CRF.model.BiLSTM_CRF import NERLSTM_CRF
from P03_NER.LSTM_CRF.utils.data_loader import word2index

config = Config()
# todo 实例化模型
models = {
    'BiLSTM': NERLSTM,
    'BiLSTM_CRF': NERLSTM_CRF
}
model = models[config.model](config.embedding_dim, config.hidden_dim, config.dropout, word2index, config.tag2id)
models_state_dict = {
    'BiLSTM': os.path.join(project_root,'P03_NER/LSTM_CRF/save_model/bilstm_best.pth'),
    'BiLSTM_CRF': os.path.join(project_root,'P03_NER/LSTM_CRF/save_model/bilstm_crf_best.pth')
}
model.load_state_dict(torch.load(models_state_dict[config.model],map_location=torch.device('cpu')))
# todo 构造id2tag：{id:tag}字典
id2tag = {v: k for k, v in config.tag2id.items()}

# todo 实现模型预测函数 model2test
def model2test(sample):
    # 构造样本训练数据  word-》id 和 掩码
    x = []
    for char in sample:
        x.append(word2index.get(char, 1)) # x:[seq_len]
    x_train = torch.tensor([x]) # x_train:[batch_size,seq_len]
    mask = (x_train != 0).long() # mask: [batch_size,seq_len]
    # print(f'model-->\n{model}')
    model.eval()
    with torch.no_grad():
        if model.name == 'BiLSTM':
            outputs = model(x_train, mask)
            preds_ids = torch.argmax(outputs, dim=-1)[0]
            tags = [id2tag[i.item()] for i in preds_ids]
        elif model.name == 'BiLSTM_CRF':
            preds_ids = model(x_train, mask)
            tags = [id2tag[i] for i in preds_ids[0]]
        chars = [i for i in sample]
        assert len(chars) == len(tags)
        result = extract_entities(chars, tags)
        return model.name, result

def extract_entities(tokens, labels):
    entities = {}  # 存储元素为（实体,实体类型）的列表
    entity, entity_type = [], None  # entity：实体字符列表，entity_type：实体类型

    for token, label in zip(tokens, labels):
        if label.startswith("B-"):  # 实体的开始
            if not entity:  # 如果entity为空
                entity.append(token)
                entity_type = label.split('-')[1]
            elif entity:  # 如果已经有实体，先把上一个识别出来实体的保存
                entities[''.join(entity)] = entity_type
                entity = []
        elif label.startswith("I-") and entity:  # 实体的中间或结尾
            entity.append(token)
        else:
            if entity:  # 保存上一个实体
                entities[''.join(entity)] = entity_type
                entity = []

    # 如果最后一个实体没有保存，手动保存
    if entity:
        entities.append((entity_type, ''.join(entity)))

    return entities

if __name__ == '__main__':
    sample = '患者精神状况好，无发热，诉右髋部疼痛，饮食差，二便正常，查体：神清，各项生命体征平稳，心肺腹查体未见异常。右髋部压痛，右下肢皮牵引固定好，无松动，右足背动脉搏动好，足趾感觉运动正常。'
    model_name, result = model2test(sample)
    print(f'{model_name}===> {result}')
